import torch
from torchvision import transforms
from PIL import Image, ImageOps

import numpy as np
import scipy.misc as misc
import os
import glob
import cv2
from utils.misc import thresh_OTSU, ReScaleSize, Crop
# from utils.model_eval import eval
from utils.evaluation_metrics import get_acc

DATABASE = './DRIVE/'
#
args = {
    # 'root'     : './dataset/' + DATABASE,
    # 'test_path': './dataset/' + DATABASE + 'test/',
    'test_path': '/home/jiayu/MyProject_2022/ROSE-2/test/',
    'pred_path': 'assets/' + 'JiaoMo/',
    'img_size': 384
}

if not os.path.exists(args['pred_path']):
    os.makedirs(args['pred_path'])


def rescale(img):
    w, h = img.size
    min_len = min(w, h)
    new_w, new_h = min_len, min_len
    scale_w = (w - new_w) // 2
    scale_h = (h - new_h) // 2
    box = (scale_w, scale_h, scale_w + new_w, scale_h + new_h)
    img = img.crop(box)
    return img


def ReScaleSize_DRIVE(image, re_size=512):
    w, h = image.size
    min_len = min(w, h)
    new_w, new_h = min_len, min_len
    scale_w = (w - new_w) // 2
    scale_h = (h - new_h) // 2
    box = (scale_w, scale_h, scale_w + new_w, scale_h + new_h)
    image = image.crop(box)
    image = image.resize((re_size, re_size))
    return image  # , origin_w, origin_h


def ReScaleSize_STARE(image, re_size=512):
    w, h = image.size
    max_len = max(w, h)
    new_w, new_h = max_len, max_len
    delta_w = new_w - w
    delta_h = new_h - h
    padding = (delta_w // 2, delta_h // 2, delta_w - (delta_w // 2), delta_h - (delta_h // 2))
    image = ImageOps.expand(image, padding, fill=0)
    # origin_w, origin_h = w, h
    image = image.resize((re_size, re_size))
    return image  # , origin_w, origin_h


def load_JiaoMo():
    test_images = []
    # test_labels = []
    for file in glob.glob(os.path.join(args['test_path'], "img", '*')):
        # basename = os.path.basename(file)
        # img_path = os.path.join(args['test_path'], 'img', basename)
        img_path = file
        # label_path = os.path.join(args['test_path'], 'gt3', basename)
        test_images.append(img_path)
        # test_labels.append(label_path)
    return test_images


def load_net():
    net = torch.load('/home/jiayu/MyProject_2022/IT_A+B+C/checkpoint_our/OCTA_16_4000.pkl')
    return net


def save_prediction(pred, filename=''):
    save_path = args['pred_path'] + 'pred/'
    if not os.path.exists(save_path):
        os.makedirs(save_path)
        print("Make dirs success!")

    mask = pred.data.cpu().numpy() * 255  # .data转换成tensor
    mask = np.transpose(np.squeeze(mask, axis=0), [1, 2, 0])
    mask = np.squeeze(mask, axis=-1)  # 只留下高和宽
    cv2.imwrite(save_path + filename + '.png', mask)


def predict():
    net = load_net()
    # images, labels = load_nerve()
    images = load_JiaoMo()
    # images, labels = load_stare()
    # images, labels = load_padova1()
    # images, labels = load_octa()

    transform = transforms.Compose([
        transforms.ToTensor()
    ])

    with torch.no_grad():
        net.eval()
        for i in range(len(images)):
            print(images[i])
            name_list = images[i].split('/')
            index = name_list[-1][:-4]  # 取去掉.tif（最后这四位）的basename
            image = Image.open(images[i]).convert("L")
            # image = cv2.imread(images[i])
            # image=image.convert("RGB")
            # label = Image.open(labels[i])
            # image, label = center_crop(image, label)

            # for other retinal vessel
            # image = rescale(image)
            # label = rescale(label)
            # image = ReScaleSize_STARE(image, re_size=args['img_size'])
            # label = ReScaleSize_DRIVE(label, re_size=args['img_size'])

            # for OCTA
            # image = Crop(image)
            # image = ReScaleSize(image)
            # label = Crop(label)
            # label = ReScaleSize(label)

            # label = label.resize((args['img_size'], args['img_size']))
            # if cuda
            image = transform(image)
            # label = transform(label)
            # image = transform(image)
            image = image.unsqueeze(0)  # 将三维提升到四维，因为网络输入要求是四维
            output = net(image)  # output.size(): [1, 1, 384, 384]
            # label.size():[1, 384, 384]
            output = output.data.cpu().numpy()
            output = output.squeeze(0)
            # output = output.transpose(2, 1, 0).squeeze(-1)
            output = output.squeeze(0)  ###预测、标签和原图方向不太一致，旋转了，要改squeeze这一块
            output = (output * 255).astype(np.uint8)
            # label = label.data.cpu().numpy().squeeze(0)
            # # label = label.transpose(2, 1, 0).squeeze(-1)
            # label = (label * 255).astype(np.uint8)

            # acc, sen, fdr = get_acc(output, label)

            # print(acc, sen, fdr)
            cv2.imwrite("/home/jiayu/MyProject_2022/ROSE-2/test/18/" + index + ".png", output)
            # cv2.imwrite("/home/imed/文档/tmp/label/" + index + ".png", label)
            # save_prediction(output, filename=index + '_pred')
            # print("output saving successfully")


if __name__ == '__main__':
    predict()
    thresh_OTSU("/home/jiayu/MyProject_2022/ROSE-2/test/18/")
